mod scale;
use anyhow::{Context, Result};
use clap::Clap;
use futures::StreamExt;
use influx_db_client::{Client, Point, Precision, Value};
use plotters::prelude::*;

#[derive(Clap, Debug)]
#[clap(version = "1.0", author = "Donggu Kang")]
struct Opts {
    #[clap(long, default_value = "ws://feed-01.ap-east-1.joomo.io:5555")]
    ping_server: String,
    #[clap(long, default_value = "http://52.7.227.223:8086")]
    influxdb: String,
    #[clap(long, default_value = "test")]
    influxdb_database: String,
    #[clap(long, default_value = "ping")]
    timeseries_name: String,
    #[clap(long)]
    debug: bool,
}

#[derive(Debug)]
struct Msg {
    rt: i64,
    ft: i64,
}

#[derive(Debug)]
struct Aggregation {
    time: i64,
    mn: i64,
    mx: i64,
}

const BATCH_SIZE: usize = 1000;

#[rocket::get("/")]
fn get_html(
    begin: rocket::State<chrono::DateTime<chrono::Utc>>,
    hist: rocket::State<std::sync::Arc<std::sync::RwLock<hdrhistogram::Histogram<u64>>>>,
) -> rocket::response::content::Html<String> {
    let hist = hist.read().unwrap();
    let end = chrono::Utc::now();

    let mut svg = String::new();
    {
        let root = SVGBackend::with_string(&mut svg, (800, 600)).into_drawing_area();
        root.fill(&WHITE).unwrap();

        let mut chart = ChartBuilder::on(&root)
            .x_label_area_size(35)
            .y_label_area_size(100)
            .margin(20)
            .caption(
                format!(
                    "Ping ft-rt ({} - {})",
                    begin.format("%Y-%m-%d %H:%M"),
                    end.format("%Y-%m-%d %H:%M")
                ),
                ("monospace", 20.0).into_font(),
            )
            .build_cartesian_2d(
                scale::InverseLogRange((0.1f64..0.9999))
                    .with_key_points(vec![0.1, 0.5, 0.9, 0.99, 0.999, 0.9999]),
                (1f64..1_000f64).log_scale(),
            )
            .unwrap();

        chart
            .configure_mesh()
            .disable_x_mesh()
            .bold_line_style(&WHITE.mix(0.3))
            .y_desc("Latency")
            .x_desc("Quantile")
            .axis_desc_style(("monospace", 15).into_font())
            .x_label_formatter(&|x| format!("{}%", x * 100.0))
            .y_label_formatter(&|y| format!("{} ms", y))
            .draw()
            .unwrap();

        chart
            .draw_series(LineSeries::new(
                hist.iter_quantiles(3).map(|it| {
                    (
                        it.quantile_iterated_to() as f64,
                        it.value_iterated_to() as f64,
                    )
                }),
                &RED,
            ))
            .unwrap();
    }
    let html = format!("<html><body>{}</body></html>", svg);
    // html
    rocket::response::content::Html(html)
}

#[tokio::main]
async fn main() -> Result<()> {
    let opts = Opts::parse();
    println!("{:?}", opts);

    // record up to 10 seconds
    let hist = std::sync::Arc::new(std::sync::RwLock::new(
        hdrhistogram::Histogram::<u64>::new_with_max(10 * 1_000_000_000, 3).unwrap(),
    ));

    let (mut tx, rx) = tokio::sync::mpsc::channel(1000);
    let mut handles = vec![];

    let (conn, _response) = tokio_tungstenite::connect_async(&opts.ping_server)
        .await
        .with_context(|| format!("fail to connect to the ping server: {}", opts.ping_server))?;

    let (_ws_writer, ws_reader) = conn.split();

    // receive ping from server
    {
        let hist = hist.clone();
        handles.push(tokio::spawn(async move {
            let mut msgs_stream = ws_reader
                .map(|msg| match msg {
                    Ok(tungstenite::Message::Binary(bytes)) => {
                        use byteorder::ReadBytesExt;
                        Msg {
                            rt: (&bytes[..])
                                .read_i64::<byteorder::BigEndian>()
                                .expect("ill-formed data received"),
                            ft: chrono::Utc::now().timestamp_nanos(),
                        }
                    }
                    _ => unreachable!(),
                })
                .chunks(BATCH_SIZE);

            while let Some(msgs) = msgs_stream.next().await {
                let mut hist = hist.write().unwrap();
                for m in &msgs {
                    let latency_ms = (m.ft - m.rt) as u64 / 1_000_000;
                    hist.record(latency_ms)?;
                }
                tx.try_send(Aggregation {
                    time: msgs[msgs.len() - 1].ft,
                    mn: msgs.iter().map(|m| m.ft - m.rt).min().unwrap(),
                    mx: msgs.iter().map(|m| m.ft - m.rt).max().unwrap(),
                })
                .context("tokio queue is full")?;
            }
            Ok::<(), anyhow::Error>(())
        }));
    }

    // send to influxdb
    let debug = opts.debug;
    handles.push(tokio::spawn(async move {
        if debug {
            println!("connecting to influxdb");
        }
        let mut client = Client::new(url::Url::parse(&opts.influxdb)?, &opts.influxdb_database)
            .set_authentication("admin", "protos123");
        client.switch_database("test");
        client.create_database(client.get_db()).await?;

        if debug {
            println!("connected to influxdb");
        }

        let mut rx = rx.map(|msg| {
            Point::new(&opts.timeseries_name)
                .add_timestamp(msg.time)
                .add_field("min", Value::Integer(msg.mn))
                .add_field("max", Value::Integer(msg.mx))
        });
        while let Some(point) = rx.next().await {
            match tokio::time::timeout(
                std::time::Duration::from_secs(1),
                client.write_point(point, Some(Precision::Nanoseconds), None),
            )
            .await
            {
                Ok(res) => {
                    res.context("fail to send to influxdb")?;
                }
                Err(_) => {
                    eprintln!("Writing to InfluxDB failed in 1 second");
                    continue;
                }
            }
        }

        println!("influx thread closed");

        Ok::<(), anyhow::Error>(())
    }));

    handles.push(tokio::spawn(async move {
        rocket::ignite()
            .manage(chrono::Utc::now())
            .manage(hist)
            .mount("/", rocket::routes![get_html])
            .launch()
            .await?;
        Ok::<(), anyhow::Error>(())
    }));

    let results = futures::future::try_join_all(handles).await?;

    let mut i = 0;
    for res in results {
        if let Err(e) = res {
            i += 1;
            eprintln!("Error {} : {:?}", i, e);
        }
    }

    Ok(())
}
